import re
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots

# 正则模式
pattern_shape = re.compile(r"M=(\d+), N=(\d+), K=(\d+)")
pattern_tflops = re.compile(r"([\d\.]+)\s+TFLOPS")

# 存储提取的数据
data = []

# 读取你的日志文件，如果不是 log.txt，请改成实际文件名
# 注意：这里我假设 './log/run_test_allshape_prehot.log' 文件存在于当前运行目录下
log_file_path = "./log/run_test_sync.log"

try:
    with open(log_file_path, "r", encoding="utf-8") as f:
        lines = f.readlines()

    # 提取数据
    for idx, line in enumerate(lines):
        m = pattern_shape.search(line)
        if m:
            current_shape = f"{m.group(1)}x{m.group(2)}x{m.group(3)}"
            for later_line in lines[idx+1:]:
                p = pattern_tflops.search(later_line)
                if p:
                    current_tflops = float(p.group(1))
                    data.append({"Shape": current_shape, "TFLOPS": current_tflops})
                    break # 找到 TFLOPS 就跳出内层循环，处理下一个 shape

except FileNotFoundError:
    print(f"错误：日志文件 '{log_file_path}' 未找到。请检查文件路径是否正确。")
    exit() # 如果文件未找到，直接退出

# 将数据转换为 Pandas DataFrame，方便 Plotly 处理
df = pd.DataFrame(data)

# 如果没有提取到数据，打印提示并退出
if df.empty:
    print("未从日志文件中提取到任何数据。请检查日志文件内容和正则表达式是否匹配。")
    exit()

# 添加一个 'Test Index' 列，用于绘图
df['Test Index'] = df.index

print(f"成功提取 {len(df)} 条数据。")

# --- 使用 Plotly Express 生成折线图 --- #
fig_line = px.line(
    df,
    x="Test Index",
    y="TFLOPS",
    title="GEMM Performance (Line Chart)",
    hover_data={"Shape": True, "TFLOPS": ":.2f"}, # 鼠标悬停时显示 Shape 和 TFLOPS
    labels={"Test Index": "测试索引", "TFLOPS": "性能 (TFLOPS)"}
)

# 更新布局，添加缩放滑块（Range Slider）和范围选择器（Range Selector）
fig_line.update_layout(
    xaxis=go.layout.XAxis(
        rangeslider=dict(visible=True), # 底部缩放滑块
        rangeselector=dict( # 快速选择缩放范围的按钮
            buttons=list([
                dict(count=1, label="1m", step="month", stepmode="backward"),
                dict(count=6, label="6m", step="month", stepmode="backward"),
                dict(count=1, label="YTD", step="year", stepmode="todate"),
                dict(count=1, label="1y", step="year", stepmode="backward"),
                dict(step="all")
            ])
        )
    ),
    # 工具提示模板，自定义显示内容
    hoverlabel=dict(
        bgcolor="white",
        font_size=12,
        font_family="sans-serif"
    )
)

# 将折线图保存为 HTML 文件
# 您只需用浏览器打开这个 HTML 文件即可享受交互式功能
output_line_html = "./html/gemm_performance_line_interactive_sync.html"
fig_line.write_html(output_line_html)
print(f"生成的交互式折线图：{output_line_html}")


# --- 使用 Plotly Express 生成柱状图 --- #
fig_bar = px.bar(
    df,
    x="Test Index",
    y="TFLOPS",
    title="GEMM Performance (Bar Chart)",
    hover_data={"Shape": True, "TFLOPS": ":.2f"}, # 鼠标悬停时显示 Shape 和 TFLOPS
    labels={"Test Index": "测试索引", "TFLOPS": "性能 (TFLOPS)"}
)

# 更新布局，添加缩放滑块和范围选择器
fig_bar.update_layout(
    xaxis=go.layout.XAxis(
        rangeslider=dict(visible=True),
        rangeselector=dict(
            buttons=list([
                dict(count=1, label="1m", step="month", stepmode="backward"),
                dict(count=6, label="6m", step="month", stepmode="backward"),
                dict(count=1, label="YTD", step="year", stepmode="todate"),
                dict(count=1, label="1y", step="year", stepmode="backward"),
                dict(step="all")
            ])
        )
    ),
    hoverlabel=dict(
        bgcolor="white",
        font_size=12,
        font_family="sans-serif"
    )
)

# 将柱状图保存为 HTML 文件
output_bar_html = "./html/gemm_performance_bar_interactive.html"
fig_bar.write_html(output_bar_html)
print(f"生成的交互式柱状图：{output_bar_html}")